# FastAPI 做backend
# FaSTapi 的功能：响应HTTP的一些协议： GET POST HEAD OPTIONS PUT PATCH DELETE（接口）
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from process import preprocess

import torch
import pickle

objects = {}
# 设置FastAPI 入口
app = FastAPI()

# 解决cors部分
# fixed （放在华为云上有风险，适用于自己电脑运行）
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 每个接口的Input
#JSON格式
class SamplePayload(BaseModel):
    function1: str
    function2: str


@app.on_event('startup') # 初始化事件（event）
# 事件的响应函数
def startup_event():
    objects['model'] = torch.load(
        'model_epoch_40.pkl', map_location=torch.device('cpu'))
    with open('__global_labels.pkl', 'rb') as fd:
        objects['global_labels'] = pickle.load(fd)
     


@app.post('/api/predict') # post是FastAPI的接口
# 从前端给后端一些输入，会用到post
def predict(payload: SamplePayload):
    #assert payload.function1 == payload.function2
    data = preprocess(payload.function1, payload.function2, objects['global_labels'])
    model = objects['model']
    prediction = torch.sigmoid(model(data)).squeeze().detach()
    print(prediction.item())
    return {"probability": prediction.item()}

